# -*- coding: utf-8 -*-
# @time: 2024/6/18 10:52
# @file: index_backtest_equal_length
# @author: tyshixi08

from get_data.origin_data import *

# 计算MACD
def MACD(df_basic, short_period=5, long_period=10, DEA_period=5):
    '''
    :param df_basic: 原始因子数据
    :param short_period: 短线滚动天数
    :param long_period: 长线滚动天数
    :param DEA_period: DEA滚动天数
    :return: dataframe数据
    '''

    df = df_basic.copy()

    data = df.iloc[:, 1]

    # 计算快线和慢线
    ema_short = data.ewm(span=short_period).mean()
    ema_long = data.ewm(span=long_period).mean()

    # 计算MACD和信号线
    DIF = ema_short - ema_long
    DEA = DIF.ewm(span=DEA_period).mean()

    MACD = 2 * (DIF - DEA)

    macd_data = pd.DataFrame({'DIF': DIF, 'DEA': DEA, 'MACD':MACD})
    df = pd.concat([df, macd_data], axis=1)
    df['signal'] = (df['MACD'] > 0).astype(int)
    return df

# 计算SMA
def SMA(df_basic, N1=5, N2=10):
    '''
    :param df_basic: 因子原始数据
    :param N1: 短线滚动天数
    :param N2: 长线滚动天数
    :return: dataframe数据
    '''

    df = df_basic.copy()

    df['date'] = pd.to_datetime(df['date'])
    df = df.set_index('date')

    fast_line = df.iloc[:, 0].rolling(window=N1).mean().reset_index().rename(columns={df.columns[0]:f'fast_line_{N1}D'})
    slow_line = df.iloc[:, 0].rolling(window=N2).mean().reset_index().rename(columns={df.columns[0]:f'slow_line_{N2}D'})
    df_line = pd.merge(fast_line, slow_line, how = 'outer', on = 'date')
    df = pd.merge(df, df_line, how = 'outer', on = 'date').dropna()
    df['signal'] = (df[f'fast_line_{N1}D'] > df[f'slow_line_{N2}D']).astype(int)
    return df

# 计算八均线
def eight_equal_length(df_basic, n1=5, n2=10, n3=21, n4=42, n5=63, n6=84, n7=126, n8=252):
    '''
    :param df_basic: 因子原始数据
    :param n1: 均线滚动天数
    :param n2: 均线滚动天数
    :param n3: 均线滚动天数
    :param n4: 均线滚动天数
    :param n5: 均线滚动天数
    :param n6: 均线滚动天数
    :param n7: 均线滚动天数
    :param n8: 均线滚动天数
    :return: dataframe数据
    '''

    df = df_basic.copy()
    df['date'] = pd.to_datetime(df['date'])

    df = df.set_index('date')
    line1 = df.iloc[:, 0].rolling(window=n1).mean().reset_index().rename(columns={df.columns[0]: f'MA_{n1}D'})
    line2 = df.iloc[:, 0].rolling(window=n2).mean().reset_index().rename(columns={df.columns[0]: f'MA_{n2}D'})
    line3 = df.iloc[:, 0].rolling(window=n3).mean().reset_index().rename(columns={df.columns[0]: f'MA_{n3}D'})
    line4 = df.iloc[:, 0].rolling(window=n4).mean().reset_index().rename(columns={df.columns[0]: f'MA_{n4}D'})
    line5 = df.iloc[:, 0].rolling(window=n5).mean().reset_index().rename(columns={df.columns[0]: f'MA_{n5}D'})
    line6 = df.iloc[:, 0].rolling(window=n6).mean().reset_index().rename(columns={df.columns[0]: f'MA_{n6}D'})
    line7 = df.iloc[:, 0].rolling(window=n7).mean().reset_index().rename(columns={df.columns[0]: f'MA_{n7}D'})
    line8 = df.iloc[:, 0].rolling(window=n8).mean().reset_index().rename(columns={df.columns[0]: f'MA_{n8}D'})
    df_line = pd.merge(line1, line2, how = 'outer', on = 'date')
    df_line = pd.merge(df_line, line3, how='outer', on='date')
    df_line = pd.merge(df_line, line4, how='outer', on='date')
    df_line = pd.merge(df_line, line5, how='outer', on='date')
    df_line = pd.merge(df_line, line6, how='outer', on='date')
    df_line = pd.merge(df_line, line7, how='outer', on='date')
    df_line = pd.merge(df_line, line8, how='outer', on='date')
    df = pd.merge(df, df_line, how='outer', on='date').dropna()

    # 计算大于八均线的数量
    ls = [n1, n2, n3, n4, n5, n6, n7, n8]
    for n in ls:
        df[f'higher_than_MA{n}D'] = (df.iloc[:, 1] > df[f'MA_{n}D']).astype(int)
    df['number_higher_than_MA'] = df.iloc[:, -8:].sum(axis=1)
    for n in ls:
        df = df.drop(f'higher_than_MA{n}D', axis=1)
    df['signal'] = (df['number_higher_than_MA'] > 5).astype(int)

    df = df.reset_index(drop=True)

    return df

# 多头排列判断
def bull_arrange(df_basic, N1=5, N2=10, N3=22):
    '''
    :param df_basic: 因子原始数据
    :param N1: 均线滚动天数
    :param N2: 均线滚动天数
    :param N3: 均线滚动天数
    :return: dataframe数据
    '''

    df = df_basic.copy()

    df['date'] = pd.to_datetime(df['date'])

    df.set_index('date', inplace=True)
    line1 = df.iloc[:, 0].rolling(window=N1).mean().reset_index().rename(columns={df.columns[0]: f'MA_{N1}D'})
    line2 = df.iloc[:, 0].rolling(window=N2).mean().reset_index().rename(columns={df.columns[0]: f'MA_{N2}D'})
    line3 = df.iloc[:, 0].rolling(window=N3).mean().reset_index().rename(columns={df.columns[0]: f'MA_{N3}D'})

    df = pd.merge(df.reset_index(), line1, how = 'outer', on = 'date')
    df = pd.merge(df, line2, how = 'outer', on = 'date')
    df = pd.merge(df, line3, how='outer', on='date')
    df = df.dropna().reset_index(drop=True)

    df['signal'] = np.nan

    for i in range(1, len(df)):
        if df.loc[i, f'MA_{N1}D'] > df.loc[i, f'MA_{N2}D'] > df.loc[i, f'MA_{N3}D']:
            df.loc[i, 'signal'] = 1
        else:
            df.loc[i, 'signal'] = 0

    return df

# 均线回测
def index_backtest(df):
    '''
    :param df: 趋势指标数据
    :return: dataframe数据
    '''

    df_index = basic_index().reset_index(drop=True).rename(columns={'交易日期':'date', '指数净值':'basic_net_value'})
    df_index['date'] = pd.to_datetime(df_index['date']).dt.strftime('%Y-%m-%d')
    df['date'] = pd.to_datetime(df['date']).dt.strftime('%Y-%m-%d')
    df = pd.merge(df_index, df, how = 'outer', on = 'date').dropna()
    df['position'] = df['signal'].shift(1)
    df = df.dropna().reset_index(drop=True)
    df['basic_net_value'] = df['basic_net_value'] / df.iloc[0, 1]

    # 累积净值计算
    df['day_return'] = 0
    for i in range(1, len(df)):
        if df.loc[i, 'position'] == 0:
            df.loc[i, 'day_return'] = 0
        else:
            df.loc[i, 'day_return'] = df.loc[i, 'basic_net_value'] / df.loc[i - 1, 'basic_net_value'] - 1

    df['cul_net_value'] = np.nan
    df.iloc[0, -1] = 1
    for i in range(1, len(df)):
        df.loc[i, 'cul_net_value'] = df.loc[i - 1, 'cul_net_value'] * (df.loc[i, 'day_return'] + 1)

    df = df.rename(columns={'date':'交易日期', 'basic_net_value':'指数净值', 'signal':'持仓信号', 'position':'持仓状态', 'day_return':'日收益率', 'cul_net_value':'累积净值'})

    return df

